import matplotlib.pyplot as plt
from PIL import Image
import os


def get_all_file_in_dir(dir_path):
    '''
    获取目录下的所有文件
    :return:
    '''
    file_name_list = []
    for root, dirs, files in os.walk(dir_path):
        if files:
            for name in files:
                file_name = '{0}/{1}'.format(root, name).replace('\\', '/')
                file_name_list.append(file_name)
    return file_name_list


def show_train_results(input_img_list, output_img_list):
    '''
    动态展示训练结果
    :param input_img_list: 一组输入图片
    :param output_img_list: 对应的一组输出图片
    :return:
    '''
    # 声明一个图框 ----------------------------------------------
    fig = plt.figure()

    # 往图框中插入一个子图， 1行2列的大小，放在第1个位置----
    axl1 = fig.add_subplot(121)
    axl1.set_title('input img')  # 设置子图标题
    axl1.axis('off')  # 子图是否开启坐标 on/off
    # ----------------------------------------------

    # 往图框中插入一个子图， 1行2列的大小，放在第2个位置----
    axl2 = fig.add_subplot(122)
    axl2.set_title('output img')  # 设置子图标题
    axl2.axis('off')  # 子图是否开启坐标 on/off
    # ----------------------------------------------

    # ----------------------------------------------------------

    for img_tup in list(zip(input_img_list, output_img_list)):
        axl1.imshow(Image.open(img_tup[0]))  # 让子图展示ta
        axl2.imshow(Image.open(img_tup[1]))  # 让子图展示ta
        plt.pause(0.01)

    plt.show()


if __name__ == '__main__':

    # 随便整点模拟数据 -------------------------------------------------------------------------------------
    img_list = get_all_file_in_dir(r'D:\机器学习数据集\celebA（分卷形式,一起解压）\celebA\img\img_align_celeba')
    input_img_list = img_list[0: 500]
    output_img_list = img_list[500: 1000]
    # ---------------------------------------------------------------------------------------------------

    # 动态展示
    show_train_results(input_img_list, output_img_list)
